"""
需求：用线性回归演示三种拟合情况
"""
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error  # 计算均方误差
from sklearn.model_selection import train_test_split


def dm01_underfitting():
    """
    演示欠拟合情况
    """
    # 1 准备数据x y(增加上噪声)
    np.random.seed(666)
    x = np.random.uniform(-3, 3, size=100)
    # np.random.normal用来增加噪声加了点噪声，代码等同于之前学过的np.random.randn(100)
    y = 0.5 * x ** 2 + x + 2 + np.random.normal(0, 1, size=100)

    # 2 实例化线性回归模型
    estimator = LinearRegression()

    # 3 训练模型
    X = x.reshape(-1, 1)
    estimator.fit(X, y)

    # 4 模型预测
    y_predict = estimator.predict(X)

    # 5 计算均方误差
    myret = mean_squared_error(y, y_predict)
    print('myret-->', myret)

    # 6 画图
    plt.scatter(x, y)
    plt.plot(np.sort(x), y_predict[np.argsort(x)], color='r')
    plt.show()


def dm02_fitting():
    """
    演示模型正好拟合
    """
    # 1 准备数据x y(增加上噪声)
    np.random.seed(666)
    x = np.random.uniform(-3, 3, size=100)
    y = 0.5 * x ** 2 + x + 2 + np.random.normal(0, 1, size=100)
    # 2 实例化线性回归模型
    estimator = LinearRegression()

    # 3 训练模型
    """
        原始数据
        特征1
        1
        2
        3
        4

        np.hstack([X, X ** 2])
        新数据
        特征1 特征2
        1      1
        2      4
        3      9
        4      16
    """
    X = x.reshape(-1, 1)
    X2 = np.hstack([X, X ** 2])  # 数据增加二次项
    estimator.fit(X2, y)

    # 4 模型预测
    y_predict = estimator.predict(X2)
    # 5 计算均方误差
    myret = mean_squared_error(y, y_predict)
    print('myret-->', myret)

    # 6 画图
    plt.scatter(x, y)
    # 画图plot折线图时 需要对x进行排序, 取x排序后对应的y值
    plt.plot(np.sort(x), y_predict[np.argsort(x)], color='r')
    plt.show()


def dm03_overfitting():
    """
    演示模型过拟合
    """
    # 1 准备数据x y(增加上噪声)
    np.random.seed(666)
    x = np.random.uniform(-3, 3, size=100)
    y = 0.5 * x ** 2 + x + 2 + np.random.normal(0, 1, size=100)
    # 2 实例化线性回归模型
    estimator = LinearRegression()

    # 3 训练模型
    X = x.reshape(-1, 1)
    X2 = np.hstack([X, X ** 10, X ** 15, X ** 20, X ** 25, X ** 30, X ** 40])  # 数据增加多次项
    estimator.fit(X2, y)

    # 4 模型预测
    y_predict = estimator.predict(X2)
    # 5 计算均方误差
    myret = mean_squared_error(y, y_predict)
    print('myret-->', myret)

    # 6 画图
    plt.scatter(x, y)
    # 画图plot折线图时 需要对x进行排序, 取x排序后对应的y值
    plt.plot(np.sort(x), y_predict[np.argsort(x)], color='r')
    plt.show()


if __name__ == '__main__':
    dm01_underfitting()
    dm02_fitting()
    dm03_overfitting()